import torch
from torchvision import datasets, transforms
import os
import argparse
import matplotlib.pyplot as plt

def get_loader(args):
    if args.dataset == 'cifar10':
        train_transform = transforms.Compose([transforms.Resize([args.image_size, args.image_size]),
                                            transforms.RandomCrop(args.image_size, padding=4),
                                            transforms.RandomHorizontalFlip(),
                                            transforms.RandAugment(),  # RandAugment augmentation for strong regularization
                                            transforms.ToTensor(),
                                            transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2470, 0.2435, 0.2616])])
        train = datasets.CIFAR10(os.path.join(args.data_path, args.dataset), train=True, download=True, transform=train_transform)
        test_transform = transforms.Compose([transforms.Resize([args.image_size, args.image_size]), transforms.ToTensor(), transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2470, 0.2435, 0.2616])])
        test = datasets.CIFAR10(os.path.join(args.data_path, args.dataset), train=False, download=True, transform=test_transform)
    else:
        print("Unknown dataset")
        exit(0)
    # Define dataloaders
    train_loader = torch.utils.data.DataLoader(dataset=train,
                                                 batch_size=args.batch_size,
                                                 shuffle=True,
                                                 num_workers=args.n_workers,
                                                 drop_last=True)

    test_loader = torch.utils.data.DataLoader(dataset=test,
                                                batch_size=args.batch_size,
                                                shuffle=False,
                                                num_workers=args.n_workers,
                                                drop_last=False)

    return train_loader, test_loader

if __name__ == '__main__':

    transform = transforms.Compose([
        transforms.ToTensor()
    ])
    train = datasets.CIFAR10('./data/cifar10', train=True, download=True,transform=transform)
    train_loader = torch.utils.data.DataLoader(dataset=train,
                                               batch_size=128,
                                               shuffle=True,
                                               num_workers=4,
                                               drop_last=True)
    img, label = next(iter(train_loader))
    for i in range(5):
        plt.figure(figsize=(8, 8))
        plt.imshow(img[i].permute(1, 2, 0))
        classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
        plt.title(f"Label: {classes[label[i]]}")
        plt.show()
    pass

